/////////////////////////////////////////////////////////////////////////////
// 2023 -
//open source under Apache License Version 2.0
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <kmhdfs/client/block_reader.h>
#include <kmhdfs/common/checksum.h>
#include "DataTransferProtocol.h"
#include <memory>
#include <kmhdfs/network/buffered_socket_reader.h>
#include <kmhdfs/network/tcp_socket.h>
#include <kmhdfs/client/packet_header.h>
#include <kmhdfs/client/peer_cache.h>
#include <kmhdfs/server/datanode_info.h>
#include <kmhdfs/server/located_blocks.h>
#include <kmhdfs/common/session_config.h>
#include <kmhdfs/client/striped_block_util.h>
#include <kmhdfs/client/thread_pool.h>
#include <kmhdfs/client/ec_chunk.h>
#include <kmhdfs/client/raw_erasure_decoder.h>

#include <future>
#include <algorithm>


namespace Hdfs {
namespace Internal {
    
class StripedInputStreamImpl;
class ReaderStrategy {
public:
    ReaderStrategy(ByteBuffer & buf);
    ReaderStrategy(ByteBuffer & buf, int length);
    ~ReaderStrategy();

    ByteBuffer & readBuf;
    int targetLength;

public:
    ByteBuffer & getReadBuf();
    int getTargetLength();
};

class StripeReader {
public:
    class CorruptedBlocks;
    class BlockReaderInfo;
    StripeReader(StripedBlockUtil::AlignedStripe & alignedStripe,
                 std::vector<LocatedBlock> & targetBlocks,
                 std::vector<BlockReaderInfo *> & readerInfos,
                 std::shared_ptr<SessionConfig> c);
    StripeReader(StripedBlockUtil::AlignedStripe & alignedStripe,
                 std::shared_ptr<ECPolicy> ecPolicy,
                 std::vector<LocatedBlock> & targetBlocks,
                 std::vector<BlockReaderInfo *> & readerInfos,
                 std::shared_ptr<CorruptedBlocks> corruptedBlocks,
                 std::shared_ptr<RawErasureDecoder> decoder,
                 StripedInputStreamImpl * dfsStripedInputStream,
                 std::shared_ptr<SessionConfig> c);

    virtual ~StripeReader();

    /**
     * read the whole stripe. do decoding if necessary
     */
    void readStripe();

    class ReaderRetryPolicy {
    public:
        ReaderRetryPolicy() {}
        ~ReaderRetryPolicy() {}

    };

    class BlockReaderInfo {
    public:
        BlockReaderInfo();
        BlockReaderInfo(std::shared_ptr<BlockReader> reader, DatanodeInfo node, long offset);
        virtual ~BlockReaderInfo() {}

        std::shared_ptr<BlockReader> reader;
        DatanodeInfo datanode;

        long blockReaderOffset;
        bool shouldSkip = false;

        void setOffset(int64_t offset);
        void skip();
    };

private:
    /**
     * Prepare all the data chunks.
     */
    virtual void prepareDecodeInputs() = 0;

    /**
     * Prepare the parity chunk and block reader if necessary.
     */
    virtual bool prepareParityChunk(int index) = 0;

    /**
     * Decode to get the missing data.
     */
    virtual void decode() = 0;

    /**
     * Default close do nothing.
     */
    virtual void close() = 0;

    void updateState4SuccessRead(StripedBlockUtil::StripingChunkReadResult & result);

    void checkMissingBlocks();

    /**
     * We need decoding. Thus go through all the data chunks and make sure we
     * submit read requests for all of them.
     */
    void readDataForDecoding();

    void readParityChunks(int num);

    std::vector<std::shared_ptr<ReaderStrategy>> getReadStrategies(StripedBlockUtil::StripingChunk & chunk);

    int32_t readToBuffer(const LocatedBlock & curBlock,
                         std::shared_ptr<BlockReader> blockReader,
                         std::shared_ptr<ReaderStrategy> strategy,
                         const DatanodeInfo & currentNode);

    int32_t readFromBlock(const LocatedBlock & curBlock,
                          const DatanodeInfo & currentNode,
                          std::shared_ptr<BlockReader> blockReader,
                          char * buf, int32_t size);

    bool readChunk(LocatedBlock& block, int chunkIndex);

    /**
     * Prepare erased indices.
     */
    std::vector<int> prepareErasedIndices();

    void clearFutures();
    
public:
    /**
     * Some fetched StripingChunk might be stored in original application
     * buffer instead of prepared decode input buffers. Some others are beyond
     * the range of the internal blocks and should correspond to all zero bytes.
     * When all pending requests have returned, this method should be called to
     * finalize decode input buffers.
     */
    void finalizeDecodeInputs();

    /**
     * Decode based on the given input buffers and erasure coding policy.
     */
    void decodeAndFillBuffer(bool fillBuffer);

protected:
    StripedBlockUtil::AlignedStripe & alignedStripe;
    std::vector<LocatedBlock> & targetBlocks;
    std::shared_ptr<CorruptedBlocks> corruptedBlocks;
    std::vector<BlockReaderInfo *> & readerInfos;
    std::shared_ptr<ECPolicy> ecPolicy;
    std::shared_ptr<RawErasureDecoder> decoder;
    StripedInputStreamImpl * dfsStripedInputStream;
    std::vector<std::shared_ptr<ECChunk>> decodeInputs;
    ThreadPool & threadPool;
    int8_t dataBlkNum;
    int8_t parityBlkNum;
    int cellSize;

private:
    std::map<int, std::future<StripedBlockUtil::BlockReadStats>> futures;
};


}
}
